// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <numeric>

#include <nebula/compute/row/compare_internal.h>
#include <nebula/testing/generator.h>
#include <nebula/testing/ktest_util.h>
#include <nebula/testing/random.h>
#include <nebula/bits/bitmap_ops.h>

namespace nebula {
namespace compute {

using nebula::bit_util::BytesForBits;
using nebula::internal::CpuInfo;
using nebula::random::RandomArrayGenerator;
using nebula::util::MiniBatch;
using nebula::util::TempVectorStack;

// Specialized case for GH-39577.
TEST(KeyCompare, CompareColumnsToRowsCuriousFSB) {
  int fsb_length = 9;
  int num_rows = 7;

  MemoryPool* pool = default_memory_pool();
  TempVectorStack stack;
  ASSERT_OK(stack.init(pool, KeyCompare::CompareColumnsToRowsTempStackUsage(num_rows)));

  auto column_right = assert_array_from_json(fixed_size_binary(fsb_length), R"([
      "000000000",
      "111111111",
      "222222222",
      "333333333",
      "444444444",
      "555555555",
      "666666666"])");
  ExecBatch batch_right({column_right}, num_rows);

  std::vector<KeyColumnMetadata> column_metadatas_right;
  ASSERT_OK(ColumnMetadatasFromExecBatch(batch_right, &column_metadatas_right));

  RowTableMetadata table_metadata_right;
  table_metadata_right.FromColumnMetadataVector(column_metadatas_right, sizeof(uint64_t),
                                                sizeof(uint64_t));

  std::vector<KeyColumnArray> column_arrays_right;
  ASSERT_OK(ColumnArraysFromExecBatch(batch_right, &column_arrays_right));

  RowTableImpl row_table;
  ASSERT_OK(row_table.init(pool, table_metadata_right));

  RowTableEncoder row_encoder;
  row_encoder.init(column_metadatas_right, sizeof(uint64_t), sizeof(uint64_t));
  row_encoder.PrepareEncodeSelected(0, num_rows, column_arrays_right);

  std::vector<uint16_t> row_ids_right(num_rows);
  std::iota(row_ids_right.begin(), row_ids_right.end(), 0);
  ASSERT_OK(row_encoder.EncodeSelected(&row_table, num_rows, row_ids_right.data()));

  auto column_left = assert_array_from_json(fixed_size_binary(fsb_length), R"([
      "000000000",
      "111111111",
      "222222222",
      "333333333",
      "444444444",
      "555555555",
      "777777777"])");
  ExecBatch batch_left({column_left}, num_rows);
  std::vector<KeyColumnArray> column_arrays_left;
  ASSERT_OK(ColumnArraysFromExecBatch(batch_left, &column_arrays_left));

  std::vector<uint32_t> row_ids_left(num_rows);
  std::iota(row_ids_left.begin(), row_ids_left.end(), 0);

  LightContext ctx{CpuInfo::GetInstance()->hardware_flags(), &stack};

  {
    uint32_t num_rows_no_match;
    std::vector<uint16_t> row_ids_out(num_rows);
    KeyCompare::CompareColumnsToRows(num_rows, nullptr, row_ids_left.data(), &ctx,
                                     &num_rows_no_match, row_ids_out.data(),
                                     column_arrays_left, row_table, true, nullptr);
    ASSERT_EQ(num_rows_no_match, 1);
    ASSERT_EQ(row_ids_out[0], 6);
  }

  {
    std::vector<uint8_t> match_bitvector(BytesForBits(num_rows));
    KeyCompare::CompareColumnsToRows(num_rows, nullptr, row_ids_left.data(), &ctx,
                                     nullptr, nullptr, column_arrays_left, row_table,
                                     true, match_bitvector.data());
    for (int i = 0; i < num_rows; ++i) {
      SCOPED_TRACE(i);
      ASSERT_EQ(nebula::bit_util::get_bit(match_bitvector.data(), i), i != 6);
    }
  }
}

// Make sure that KeyCompare::CompareColumnsToRows uses no more stack space than declared
// in KeyCompare::CompareColumnsToRowsTempStackUsage().
TEST(KeyCompare, CompareColumnsToRowsTempStackUsage) {
  for (auto num_rows :
       {0, 1, MiniBatch::kMiniBatchLength, MiniBatch::kMiniBatchLength * 64}) {
    SCOPED_TRACE("num_rows = " + std::to_string(num_rows));

    MemoryPool* pool = default_memory_pool();
    TempVectorStack stack;
    ASSERT_OK(stack.init(pool, KeyCompare::CompareColumnsToRowsTempStackUsage(num_rows)));

    RandomArrayGenerator gen(42);

    auto column_right = gen.Int8(num_rows, 0, 127);
    ExecBatch batch_right({column_right}, num_rows);

    std::vector<KeyColumnMetadata> column_metadatas_right;
    ASSERT_OK(ColumnMetadatasFromExecBatch(batch_right, &column_metadatas_right));

    RowTableMetadata table_metadata_right;
    table_metadata_right.FromColumnMetadataVector(column_metadatas_right,
                                                  sizeof(uint64_t), sizeof(uint64_t));

    std::vector<KeyColumnArray> column_arrays_right;
    ASSERT_OK(ColumnArraysFromExecBatch(batch_right, &column_arrays_right));

    RowTableImpl row_table;
    ASSERT_OK(row_table.init(pool, table_metadata_right));

    RowTableEncoder row_encoder;
    row_encoder.init(column_metadatas_right, sizeof(uint64_t), sizeof(uint64_t));
    row_encoder.PrepareEncodeSelected(0, num_rows, column_arrays_right);

    std::vector<uint16_t> row_ids_right(num_rows);
    std::iota(row_ids_right.begin(), row_ids_right.end(), 0);
    ASSERT_OK(row_encoder.EncodeSelected(&row_table, num_rows, row_ids_right.data()));

    auto column_left = gen.Int8(num_rows, 0, 127);
    ExecBatch batch_left({column_left}, num_rows);
    std::vector<KeyColumnArray> column_arrays_left;
    ASSERT_OK(ColumnArraysFromExecBatch(batch_left, &column_arrays_left));

    std::vector<uint32_t> row_ids_left(num_rows);
    std::iota(row_ids_left.begin(), row_ids_left.end(), 0);

    LightContext ctx{CpuInfo::GetInstance()->hardware_flags(), &stack};

    uint32_t num_rows_no_match;
    std::vector<uint16_t> row_ids_out(num_rows);
    KeyCompare::CompareColumnsToRows(num_rows, nullptr, row_ids_left.data(), &ctx,
                                     &num_rows_no_match, row_ids_out.data(),
                                     column_arrays_left, row_table, true, nullptr);
  }
}

// Compare columns to rows at offsets over 2GB within a row table.
// Certain AVX2 instructions may behave unexpectedly causing troubles like GH-41813.
TEST(KeyCompare, LARGE_MEMORY_TEST(CompareColumnsToRowsLarge)) {
  if constexpr (sizeof(void*) == 4) {
    KTEST_SKIP() << "Test only works on 64-bit platforms";
  }

  // The idea of this case is to create a row table using several fixed length columns and
  // one var length column (so the row is hence var length and has offset buffer), with
  // the overall data size exceeding 2GB. Then compare each row with itself.
  constexpr int64_t two_gb = 2ll * 1024ll * 1024ll * 1024ll;
  // The compare function requires the row id of the left column to be uint16_t, hence the
  // number of rows.
  constexpr int64_t num_rows = std::numeric_limits<uint16_t>::max() + 1;
  const std::vector<std::shared_ptr<DataType>> fixed_length_types{uint64(), uint32()};
  // The var length column should be a little smaller than 2GB to workaround the capacity
  // limitation in the var length builder.
  constexpr int32_t var_length = two_gb / num_rows - 1;
  auto row_size = std::accumulate(fixed_length_types.begin(), fixed_length_types.end(),
                                  static_cast<int64_t>(var_length),
                                  [](int64_t acc, const std::shared_ptr<DataType>& type) {
                                    return acc + type->byte_width();
                                  });
  // The overall size should be larger than 2GB.
  ASSERT_GT(row_size * num_rows, two_gb);

  MemoryPool* pool = default_memory_pool();

  // The left side columns.
  std::vector<KeyColumnArray> columns_left;
  ExecBatch batch_left;
  {
    std::vector<Datum> values;

    // Several fixed length arrays containing random content.
    for (const auto& type : fixed_length_types) {
      ASSERT_OK_AND_ASSIGN(auto value, ::nebula::gen::Random(type)->Generate(num_rows));
      values.push_back(std::move(value));
    }
    // A var length array containing 'X' repeated var_length times.
    ASSERT_OK_AND_ASSIGN(auto value_var_length,
                         ::nebula::gen::Constant(
                             std::make_shared<BinaryScalar>(std::string(var_length, 'X')))
                             ->Generate(num_rows));
    values.push_back(std::move(value_var_length));

    batch_left = ExecBatch(std::move(values), num_rows);
    ASSERT_OK(ColumnArraysFromExecBatch(batch_left, &columns_left));
  }

  // The right side row table.
  RowTableImpl row_table_right;
  {
    // Encode the row table with the left columns.
    std::vector<KeyColumnMetadata> column_metadatas;
    ASSERT_OK(ColumnMetadatasFromExecBatch(batch_left, &column_metadatas));
    RowTableMetadata table_metadata;
    table_metadata.FromColumnMetadataVector(column_metadatas, sizeof(uint64_t),
                                            sizeof(uint64_t));
    ASSERT_OK(row_table_right.init(pool, table_metadata));
    std::vector<uint16_t> row_ids(num_rows);
    std::iota(row_ids.begin(), row_ids.end(), 0);
    RowTableEncoder row_encoder;
    row_encoder.init(column_metadatas, sizeof(uint64_t), sizeof(uint64_t));
    row_encoder.PrepareEncodeSelected(0, num_rows, columns_left);
    ASSERT_OK(row_encoder.EncodeSelected(
        &row_table_right, static_cast<uint32_t>(num_rows), row_ids.data()));

    // The row table must contain an offset buffer.
    ASSERT_NE(row_table_right.offsets(), nullptr);
    // The whole point of this test.
    ASSERT_GT(row_table_right.offsets()[num_rows - 1], two_gb);
  }

  // The rows to compare.
  std::vector<uint32_t> row_ids_to_compare(num_rows);
  std::iota(row_ids_to_compare.begin(), row_ids_to_compare.end(), 0);

  TempVectorStack stack;
  ASSERT_OK(stack.init(pool, KeyCompare::CompareColumnsToRowsTempStackUsage(num_rows)));
  LightContext ctx{CpuInfo::GetInstance()->hardware_flags(), &stack};

  {
    // No selection, output no match row ids.
    uint32_t num_rows_no_match;
    std::vector<uint16_t> row_ids_out(num_rows);
    KeyCompare::CompareColumnsToRows(num_rows, /*sel_left_maybe_null=*/nullptr,
                                     row_ids_to_compare.data(), &ctx, &num_rows_no_match,
                                     row_ids_out.data(), columns_left, row_table_right,
                                     /*are_cols_in_encoding_order=*/true,
                                     /*out_match_bitvector_maybe_null=*/nullptr);
    ASSERT_EQ(num_rows_no_match, 0);
  }

  {
    // No selection, output match bit vector.
    std::vector<uint8_t> match_bitvector(BytesForBits(num_rows));
    KeyCompare::CompareColumnsToRows(
        num_rows, /*sel_left_maybe_null=*/nullptr, row_ids_to_compare.data(), &ctx,
        /*out_num_rows=*/nullptr, /*out_sel_left_maybe_same=*/nullptr, columns_left,
        row_table_right,
        /*are_cols_in_encoding_order=*/true, match_bitvector.data());
    ASSERT_EQ(nebula::internal::CountSetBits(match_bitvector.data(), 0, num_rows),
              num_rows);
  }

  std::vector<uint16_t> selection_left(num_rows);
  std::iota(selection_left.begin(), selection_left.end(), 0);

  {
    // With selection, output no match row ids.
    uint32_t num_rows_no_match;
    std::vector<uint16_t> row_ids_out(num_rows);
    KeyCompare::CompareColumnsToRows(num_rows, selection_left.data(),
                                     row_ids_to_compare.data(), &ctx, &num_rows_no_match,
                                     row_ids_out.data(), columns_left, row_table_right,
                                     /*are_cols_in_encoding_order=*/true,
                                     /*out_match_bitvector_maybe_null=*/nullptr);
    ASSERT_EQ(num_rows_no_match, 0);
  }

  {
    // With selection, output match bit vector.
    std::vector<uint8_t> match_bitvector(BytesForBits(num_rows));
    KeyCompare::CompareColumnsToRows(
        num_rows, selection_left.data(), row_ids_to_compare.data(), &ctx,
        /*out_num_rows=*/nullptr, /*out_sel_left_maybe_same=*/nullptr, columns_left,
        row_table_right,
        /*are_cols_in_encoding_order=*/true, match_bitvector.data());
    ASSERT_EQ(nebula::internal::CountSetBits(match_bitvector.data(), 0, num_rows),
              num_rows);
  }
}

}  // namespace compute
}  // namespace nebula
